TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
Rollout 引导的 MoE LLM 强化学习 FP4 量化感知训练框架
🧭 低精度 rollout 加速 MoE LLM 强化学习训练
Alibaba Token Hub, Alibaba Group · Ohio State University
提出 TRACE，通过 rollout 侧量化结果引导训练侧 FP4 舍入，实现 FP4 rollout 与 BF16 性能相当并最高加速 5.4 倍
📌 工作参考价值（4 点）
- 在 RL 调 LLM 推荐/搜索 Agent 或生成式 ranker 时，rollout 阶段是主要算力瓶颈，可尝试 FP4 权重/激活 + FP4 KV cache 联合做 rollout；TRACE 表明可与 BF16 rollout 获得相近的最终策略质量。
- 避免直接用离线 PTQ 量化 BF16 训练好的策略：把 rollout 侧实际量化噪声（舍入结果、scale）作为监督信号引入训练侧 QAT，能显著改善最终 FP4 策略性能，比事后量化更可靠。
- 工程上可借鉴 selective quantization-info caching：只缓存深层 mantissa/scale，降低多 GPU 或分布式 rollout worker 间的存储与通信开销，对 MoE 或超大 LLM 生成服务有直接参考。
- 低精度 rollout 最高 5.4 倍加速意味着相同算力下可大幅增加采样轨迹数量，适合需要大量探索的 query 推荐、对话式推荐 Agent 等 RL 训练场景。